package org.usmile.algorithms.learn.tree;

public class AvlTree<E extends Comparable<E>> {
    private TreeNode<E> root;

    public void add(E e) {
        root = add(root, e);
    }

    private TreeNode<E> add(TreeNode<E> node, E e) {
        if (node == null) {
            return new TreeNode<>(e);
        }

        if (e.compareTo(node.e) == 0) {
            return node;
        } else if (e.compareTo(node.e) < 0) {
            node.left = add(node.left, e);
        } else {
            node.right = add(node.right, e);
        }

        return reBalance(node);
    }

    public void remove(E e) {
        if (root == null) {
            return;
        }

        root = remove(root, e);
    }

    private TreeNode<E> remove(TreeNode<E> node, E e) {
        if (node == null) {
            return null;
        }

        if (e.compareTo(node.e) > 0) {
            node.right = remove(node.right, e);
        } else if (e.compareTo(node.e) < 0) {
            node.left = remove(node.left, e);
        } else {
            if (node.left == null) {
                node = node.right;
            } else if (node.right == null) {
                node = node.left;
            } else {
                // node 的 left 和 right 都不为空
                TreeNode<E> rightMin = getMin(node.right);
                rightMin.right = remove(node.right, rightMin.e);
                rightMin.left = node.left;

                node = rightMin;
            }
        }

        return reBalance(node);
    }

    private TreeNode<E> getMin(TreeNode<E> node) {
        if (node.left == null) {
            return node;
        }

        return getMin(node.left);
    }

    private TreeNode<E> reBalance(TreeNode<E> node) {
        if (node == null) {
            return null;
        }
        node.height = Math.max(getHeight(node.left), getHeight(node.right)) + 1;
        int balanceFactor = getBalanceFactor(node);
        if (balanceFactor > 1) {
            if (getBalanceFactor(node.left) < 0) {
                node.left = leftRotate(node.left); // LR
            }

            return rightRotate(node); // LL
        } else if (balanceFactor < -1) {
            if (getBalanceFactor(node.right) > 0) {
                node.right = rightRotate(node.right); // RL
            }

            return leftRotate(node); // RR
        }

        return node;
    }


    // 对节点 y 进行向右旋转操作，返回旋转后新的根节点 x
    //        y                                    x
    //       / \                                 /   \
    //      x   T4        向右旋转 (y)          z       y
    //     / \          --------------->      / \     / \
    //    z   T3                            T1   T2 T3   T4
    //   / \
    // T1   T2
    private TreeNode<E> rightRotate(TreeNode<E> node) {
        TreeNode<E> left = node.left;
        TreeNode<E> leftRight = left.right;
        left.right = node;
        node.left = leftRight;

        node.height = Math.max(getHeight(node.left), getHeight(node.right)) + 1;
        left.height = Math.max(getHeight(left.left), getHeight(left.right)) + 1;

        return left;
    }

    // 对节点y进行向左旋转操作，返回旋转后新的根节点x
    //    y                             x
    //  /  \                          /   \
    // T4   x      向左旋转 (y)       y     z
    //     / \   - - - - - - - ->   / \   / \
    //   T3  z                     T4 T3 T1 T2
    //      / \
    //     T1 T2
    private TreeNode<E> leftRotate(TreeNode<E> node) {
        TreeNode<E> right = node.right;
        TreeNode<E> rightLeft = right.left;
        right.left = node;
        node.right = rightLeft;

        node.height = Math.max(getHeight(node.left), getHeight(node.right)) + 1;
        right.height = Math.max(getHeight(right.left), getHeight(right.right)) + 1;

        return right;
    }

    public boolean isBalance() {
        return isBalance(root);
    }

    private boolean isBalance(TreeNode<E> node) {
        if (node == null) {
            return true;
        }

        if (Math.abs(getBalanceFactor(node)) > 1) {
            return false;
        }

        return isBalance(node.left) && isBalance(node.right);
    }

    private int getHeight(TreeNode<E> node) {
        if (node == null) {
            return 0;
        }

        return node.height;
    }

    private int getBalanceFactor(TreeNode<E> node) {
        if (node == null) {
            return 0;
        }

        return getHeight(node.left) - getHeight(node.right);
    }

    public static void main(String[] args) {
        AvlTree<Integer> tree = new AvlTree<>();
        tree.add(1);
        tree.add(2);
        tree.add(3);
        tree.add(4);
        tree.add(5);
        tree.add(6);
        tree.add(-1);
        tree.add(-2);
        tree.add(-3);
        tree.add(-4);
        tree.add(-5);
        tree.add(-6);
        TreeUtil.print(tree.root);
        tree.remove(1);
        TreeUtil.print(tree.root);
        tree.remove(2);
        TreeUtil.print(tree.root);
        tree.remove(3);
        TreeUtil.print(tree.root);
        tree.remove(-6);
        TreeUtil.print(tree.root);
        tree.remove(-5);
        TreeUtil.print(tree.root);
        tree.add(-5);
        TreeUtil.print(tree.root);
    }
}
